import onnx
from onnx import helper
from onnx import TensorProto
import numpy as np

inp = helper.make_tensor_value_info('X', TensorProto.FLOAT, [1,6,3,6])
# weight = np.random.randn(108)
# const_input = helper.make_tensor('X', TensorProto.FLOAT, [1,6,3,6], weight)


shape = helper.make_tensor_value_info('shape', TensorProto.INT64, [4])
shape_weight = np.array([1, 3, 6, 6])
const_shape = helper.make_tensor('shape', TensorProto.INT64, [4], shape_weight)

outp = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [1,3,6,6])
reshape_def = helper.make_node(
    'Reshape',
    ['X', 'shape'],
    ['Y']
)

conv_weight = np.random.randn(27)
W = helper.make_tensor('W', TensorProto.FLOAT, [1, 3, 3, 3], conv_weight)
B = helper.make_tensor('B', TensorProto.FLOAT, [1,1,1,1], [0.54964069])
YC = helper.make_tensor_value_info('YC', TensorProto.FLOAT, [1, 1, 4, 4])

conv_def = helper.make_node(
    'Conv', # node name
    ['Y', 'W', 'B'],
    ['YC'], # outputs
    # attributes
    strides=[1,1],
    )

graph_def = helper.make_graph(
    [reshape_def, conv_def],
    "test_reshape_model",
    [inp, shape],
    [YC],
    initializer=[const_shape, W, B]
)

mode_def = helper.make_model(graph_def, producer_name='onnx-example')
onnx.checker.check_model(mode_def)
onnx.save(mode_def, "./reshape.onnx")